/*
 * @Author: wuqingchun
 * @Date: 2024-05-22 17:17:53
 * @Description: 协程类实现
 * @LastEditors: wuqingchun
 * @LastEditTime: 2024-05-24 09:33:26
 */
#ifndef __COROUTINE_H__
#define __COROUTINE_H__

#include <iostream>

#include <ucontext.h>

#include <stdio.h>
#include <string.h>
#include <errno.h>
#include <array>

enum class CoState
{
    CO_STATE_INIT,
    CO_STATE_RUNNING, // 正在运行
    CO_STATE_SUSPEND, // 被挂起
};

struct CoContext
{
    static constexpr size_t COROUTINE_STACK_SIZE = 16 << 12;
    static constexpr uint32_t INVALID_COROUTINE_ID = (uint32_t)-1;

    ucontext_t stUCtx;
    uint32_t uiCtxId;             // 协程ID
    CoState enState; // 协程状态
    char acStack[COROUTINE_STACK_SIZE];
    void* pUserData;
};

class Coroutine;

typedef void (*COROUTINE_FUNC)(void* pArg);
typedef void (*COROUTINE_FUNC_NOARG)(void);

class Coroutine
{
public:
    Coroutine(ucontext_t& refMainUCtx, bool bRunning):
        m_refMainUCtx(refMainUCtx),
        m_refRunning(bRunning)
    {

    }

    ~Coroutine()
    {
        if (m_pcStack != nullptr)
        {
            delete [] m_pcStack;
            m_pcStack = nullptr;
        }
    }

    int Init(size_t ulStackSize, COROUTINE_FUNC pfFunc)
    {
        char* pcStack = new (std::nothrow) char[ulStackSize];
        if (pcStack == nullptr)
        {
            return -1;
        }

        int iErr = getcontext(&m_stUCtx);
        if (iErr == -1)
        {
            printf("getcontext error, %s\n", strerror(errno));

            return -1;
        }

        m_stUCtx.uc_stack.ss_sp = pcStack;
        m_stUCtx.uc_stack.ss_size = ulStackSize;  
        m_stUCtx.uc_link = &m_refMainUCtx; // 设置返回时切换到主流程
        makecontext(&m_stUCtx, (void(*)())pfFunc, 1, this);

        m_pcStack = pcStack;

        return 0;
    }

    int Yeild()
    {
        int iErr = swapcontext(&m_stUCtx, &m_refMainUCtx); // 切换到主协程
        if (iErr == -1)
        {
            printf("swapcontext error, %s\n", strerror(errno));
        }

        return iErr;
    }

    int Resume()
    {
        int iErr = swapcontext(&m_refMainUCtx, &m_stUCtx); // 切换到当前协程
        if (iErr == -1)
        {
            printf("swapcontext error, %s\n", strerror(errno));
        }

        return iErr;
    }

    bool IsRunning()
    {
        return m_refRunning;
    }

private:
    char* m_pcStack = nullptr;
    ucontext_t m_stUCtx;
    ucontext_t& m_refMainUCtx;
    bool& m_refRunning;
};

class CoroutineManager;

typedef void (*COROUTINE_FUNC_PF)(CoroutineManager* pobjManager, uint32_t uiCoId);


class CoroutineManager
{
private:
    static constexpr size_t COROUTINE_COUNT = 1024;
    
    
public:
    CoroutineManager()
    {

    }

    ~CoroutineManager()
    {
        for (size_t ulIndex = 0; ulIndex < COROUTINE_COUNT; ulIndex++)
        {
            if (m_arrSubCo[ulIndex] != nullptr)
            {
                __Yield(*m_arrSubCo[ulIndex]);
                delete m_arrSubCo[ulIndex];
                m_arrSubCo[ulIndex] = nullptr;
            }
        }
    }

    uint32_t CoAdd(COROUTINE_FUNC_PF pfFunc)
    {
        if (m_uiSubCoCount == COROUTINE_COUNT)
        {
            return CoContext::INVALID_COROUTINE_ID;
        }

        CoContext* pstCo = new (std::nothrow) CoContext;
        if (pstCo == nullptr)
        {
            return CoContext::INVALID_COROUTINE_ID;
        }
    
        int iErr = getcontext(&pstCo->stUCtx);
        if (iErr == -1)
        {
            printf("getcontext error, %s\n", strerror(errno));
            delete pstCo;
            return CoContext::INVALID_COROUTINE_ID;
        }

        pstCo->stUCtx.uc_stack.ss_sp = pstCo->acStack;
        pstCo->stUCtx.uc_stack.ss_size = sizeof(pstCo->acStack);  
        pstCo->stUCtx.uc_link = &m_stMainUCtx; // 设置返回时切换到主流程
        pstCo->uiCtxId = m_uiNextID;
        m_arrSubCo[m_uiNextID] = pstCo;
        makecontext(&pstCo->stUCtx, (void (*)())pfFunc, 2, this, m_uiNextID);

        m_uiSubCoCount++;
        return m_uiNextID++;
    }

    void CoDelete(uint32_t uiCoId)
    {
        if (uiCoId == CoContext::INVALID_COROUTINE_ID)
        {
            return;
        }

        CoContext* pstCo = m_arrSubCo[uiCoId];
        if (pstCo == nullptr)
        {
            return;
        }

        delete pstCo;
        m_arrSubCo[uiCoId] = nullptr;
        m_uiSubCoCount--;

        if (m_uiNextID > uiCoId)
        {
            m_uiNextID = uiCoId;
        }
    }

    int CoResume(uint32_t uiCoId, void* pUserData)
    {
        if (uiCoId == CoContext::INVALID_COROUTINE_ID)
        {
            return -1;
        }

        return __Resume(*m_arrSubCo[uiCoId], pUserData);
    }

    int CoYeild(uint32_t uiCoId)
    {
        if (uiCoId == CoContext::INVALID_COROUTINE_ID)
        {
            return -1;
        }
        
        return __Yield(*m_arrSubCo[uiCoId]);
    }

    CoContext* GetCoContext(uint32_t uiCoId)
    {
        if (uiCoId == CoContext::INVALID_COROUTINE_ID)
        {
            return nullptr;
        }

        return m_arrSubCo[uiCoId];
    }
private:
    int __Yield(CoContext& refCoCtx)
    {
        if (refCoCtx.enState != CoState::CO_STATE_RUNNING)
        {
            return 0;
        }

        int iErr = swapcontext(&refCoCtx.stUCtx, &m_stMainUCtx);          // 切换到主协程
        if (iErr == -1)
        {
            printf("swapcontext error, %s\n", strerror(errno));
        }

        return iErr;
    }

    int __Resume(CoContext& refCoCtx, void* pUserData)
    {
        if (refCoCtx.enState == CoState::CO_STATE_RUNNING)
        {
            return 0;
        }
        
        refCoCtx.pUserData = pUserData;
        int iErr = swapcontext(&m_stMainUCtx, &refCoCtx.stUCtx);          // 切换到主协程
        if (iErr == -1)
        {
            printf("swapcontext error, %s\n", strerror(errno));
        }

        return iErr;
    }
private:
    ucontext_t m_stMainUCtx;                                       // 主协程
    std::array<CoContext*, COROUTINE_COUNT> m_arrSubCo;     // 子协程指针数组
    uint32_t m_uiSubCoCount = 0;
    uint32_t m_uiNextID = 0;
};

#endif //__COROUTINE_H__
